#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
AI自主探索谱面规律的主脚本
采用渐进式迭代方法让AI从谱面和音频特征数据中发现规律
使用ModelScope API替代Ollama
"""

import json
import os
import random
from typing import List, Dict, Any
from openai import OpenAI

class PatternDiscoveryEngine:
    def __init__(self, data_dir: str, batch_size: int = 10):
        """
        初始化规律发现引擎
        
        Args:
            data_dir: 包含谱面和音频特征JSON文件的目录
            batch_size: 每批处理的样本数量
        """
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.batches = []
        self.accumulated_patterns = ""
        
        # 初始化ModelScope客户端
        self.client = OpenAI(
            base_url='https://api-inference.modelscope.cn/v1',
            api_key='ms-4611f4fa-6ec0-4ba7-87da-43032bc42e49'  # ModelScope Token
        )
        
    def organize_data(self):
        """组织数据成批次"""
        # 查找所有JSON文件
        json_files = [f for f in os.listdir(self.data_dir) if f.endswith('.json')]
        
        # 打乱文件顺序以增加随机性
        random.shuffle(json_files)
        
        # 分成批次
        self.batches = []
        for i in range(0, len(json_files), self.batch_size):
            batch = json_files[i:i+self.batch_size]
            self.batches.append(batch)
        
        print(f"总共找到 {len(json_files)} 个JSON文件，分为 {len(self.batches)} 批次")
        
    def load_batch_data(self, batch_files: List[str]) -> List[Dict]:
        """加载一批数据"""
        batch_data = []
        for file in batch_files:
            file_path = os.path.join(self.data_dir, file)
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    batch_data.append(data)
            except Exception as e:
                print(f"加载文件 {file} 时出错: {e}")
        return batch_data
    
    def discover_patterns_with_ai(self, batch_data: List[Dict], previous_patterns: str = "") -> str:
        """
        使用ModelScope API从一批数据中发现规律
        
        Args:
            batch_data: 一批谱面和音频特征数据
            previous_patterns: 之前发现的规律
            
        Returns:
            新发现的规律
        """
        prompt = f"""
你是一位音乐游戏谱面分析专家。{'之前我们已经发现了以下规律：' + previous_patterns if previous_patterns else '请分析以下谱面数据，发现其中的规律'}

请分析这些谱面与其音频特征的关系，重点关注：
1. 音频特征如何影响谱面难度设计
2. 音符分布与音乐节拍的对应关系
3. 不同难度等级谱面的结构特点
4. 谱面密度与音频能量的关系
5. 特殊音符类型（如Hold、Slide）的使用规律

数据：
{json.dumps(batch_data, ensure_ascii=False, indent=2)}

请以清晰、结构化的方式总结你发现的规律。
"""
        
        try:
            response = self.client.chat.completions.create(
                model='Qwen/Qwen3-Next-80B-A3B-Instruct',
                messages=[
                    {
                        'role': 'system',
                        'content': '你是一位音乐游戏谱面分析专家，擅长从数据中发现规律。'
                    },
                    {
                        'role': 'user',
                        'content': prompt
                    }
                ],
                temperature=0.7,
                max_tokens=2000
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"调用ModelScope API时出错: {e}")
            return "未能成功获取规律"
    
    def consolidate_patterns(self, old_patterns: str, new_patterns: str) -> str:
        """
        整合新旧规律
        
        Args:
            old_patterns: 旧规律
            new_patterns: 新规律
            
        Returns:
            整合后的规律
        """
        if not old_patterns:
            return new_patterns
            
        consolidation_prompt = f"""
请整合以下两组规律发现：
第一组（已有规律）：
{old_patterns}

第二组（新增规律）：
{new_patterns}

请合并、去重并完善这些规律，形成更全面的总结。要求：
1. 保留所有有价值的信息
2. 去除重复内容
3. 以清晰的结构组织规律
4. 突出最重要的发现
"""
        
        try:
            response = self.client.chat.completions.create(
                model='Qwen/Qwen3-Next-80B-A3B-Instruct',
                messages=[
                    {
                        'role': 'system',
                        'content': '你是一位专业的信息整合专家，擅长整理和归纳信息。'
                    },
                    {
                        'role': 'user',
                        'content': consolidation_prompt
                    }
                ],
                temperature=0.5,
                max_tokens=2000
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"整合规律时出错: {e}")
            return old_patterns + "\n" + new_patterns
    
    def validate_patterns(self, sample_data: Dict, patterns: str) -> str:
        """
        验证规律的准确性
        
        Args:
            sample_data: 样本数据
            patterns: 当前规律
            
        Returns:
            验证结果和改进建议
        """
        validation_prompt = f"""
作为音乐游戏谱面专家，请基于以下规律为给定音频特征生成谱面，然后评估生成结果的合理性：

规律：
{patterns}

音频特征：
{json.dumps(sample_data.get('features', {}), ensure_ascii=False, indent=2)}

难度等级：{sample_data.get('difficulty', 'unknown')}

请按以下步骤进行：
1. 基于规律简要说明你会如何生成谱面
2. 与实际谱面对比，评估准确性和合理性
3. 如果发现问题，请指出并提出改进建议
"""
        
        try:
            response = self.client.chat.completions.create(
                model='Qwen/Qwen3-Next-80B-A3B-Instruct',
                messages=[
                    {
                        'role': 'system',
                        'content': '你是一位音乐游戏谱面验证专家，擅长评估谱面生成质量。'
                    },
                    {
                        'role': 'user',
                        'content': validation_prompt
                    }
                ],
                temperature=0.5,
                max_tokens=1500
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"验证规律时出错: {e}")
            return "验证失败"
    
    def refine_patterns(self, patterns: str, feedback: str) -> str:
        """
        根据反馈精炼规律
        
        Args:
            patterns: 当前规律
            feedback: 反馈意见
            
        Returns:
            精炼后的规律
        """
        refinement_prompt = f"""
请根据以下反馈精炼和完善规律：

当前规律：
{patterns}

反馈意见：
{feedback}

请修改规律，使其更加准确和完善。要求：
1. 保留有价值的内容
2. 根据反馈修正错误或不准确的部分
3. 保持规律的清晰性和实用性
"""
        
        try:
            response = self.client.chat.completions.create(
                model='Qwen/Qwen3-Next-80B-A3B-Instruct',
                messages=[
                    {
                        'role': 'system',
                        'content': '你是一位专业的知识提炼专家，擅长优化和完善规则体系。'
                    },
                    {
                        'role': 'user',
                        'content': refinement_prompt
                    }
                ],
                temperature=0.5,
                max_tokens=2000
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"精炼规律时出错: {e}")
            return patterns
    
    def run_discovery_iteration(self) -> str:
        """
        运行一轮规律发现
        
        Returns:
            本轮发现的规律
        """
        print("开始新一轮规律发现...")
        
        for i, batch_files in enumerate(self.batches):
            print(f"  处理第 {i+1}/{len(self.batches)} 批数据...")
            batch_data = self.load_batch_data(batch_files)
            
            if not batch_data:
                continue
                
            new_patterns = self.discover_patterns_with_ai(batch_data, self.accumulated_patterns)
            self.accumulated_patterns = self.consolidate_patterns(self.accumulated_patterns, new_patterns)
            
            # 每处理10批数据就做一个阶段性总结
            if (i + 1) % 10 == 0:
                print(f"  完成前 {i+1} 批数据的阶段性规律总结")
        
        return self.accumulated_patterns
    
    def run_validation_phase(self) -> str:
        """
        运行验证阶段
        
        Returns:
            验证后精炼的规律
        """
        print("开始验证阶段...")
        
        # 随机选择一些样本进行验证
        validation_samples = random.sample(
            [file for batch in self.batches for file in batch], 
            min(50, len([file for batch in self.batches for file in batch]))
        )
        
        for i, sample_file in enumerate(validation_samples):
            print(f"  验证第 {i+1}/{len(validation_samples)} 个样本...")
            file_path = os.path.join(self.data_dir, sample_file)
            
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    sample_data = json.load(f)
            except Exception as e:
                print(f"加载样本 {sample_file} 时出错: {e}")
                continue
            
            validation_result = self.validate_patterns(sample_data, self.accumulated_patterns)
            # 根据验证结果精炼规律
            self.accumulated_patterns = self.refine_patterns(self.accumulated_patterns, validation_result)
        
        return self.accumulated_patterns
    
    def discover_patterns(self, num_iterations: int = 3) -> str:
        """
        主函数：执行完整的规律发现过程
        
        Args:
            num_iterations: 迭代轮数
            
        Returns:
            最终发现的规律
        """
        print("开始AI自主探索谱面规律...")
        
        # 组织数据
        self.organize_data()
        
        if not self.batches:
            raise ValueError("未找到任何数据文件")
        
        # 多轮迭代
        for iteration in range(num_iterations):
            print(f"\n=== 第 {iteration + 1}/{num_iterations} 轮迭代 ===")
            
            # 规律发现阶段
            self.run_discovery_iteration()
            
            # 验证和精炼阶段
            self.run_validation_phase()
            
            print(f"第 {iteration + 1} 轮迭代完成")
            print(f"当前规律摘要: {self.accumulated_patterns[:200]}...")
        
        # 生成最终的规律提示词
        final_prompt = self.generate_final_prompt()
        
        return final_prompt
    
    def generate_final_prompt(self) -> str:
        """
        生成最终的规律提示词
        
        Returns:
            最终的规律提示词
        """
        final_prompt_generation = f"""
基于以上发现的所有规律，请生成一个结构化的提示词模板，用于指导AI生成舞萌DX谱面。

提示词应包含以下部分：
1. 角色设定
2. 任务说明
3. 音频特征分析方法
4. 谱面生成规则
5. 不同难度等级的处理方式
6. 特殊音符类型的使用原则
7. 输出格式要求

规律总结：
{self.accumulated_patterns}
"""
        
        try:
            response = self.client.chat.completions.create(
                model='Qwen/Qwen3-Next-80B-A3B-Instruct',
                messages=[
                    {
                        'role': 'system',
                        'content': '你是一位专业的提示词工程师，擅长构建结构化的AI指令。'
                    },
                    {
                        'role': 'user',
                        'content': final_prompt_generation
                    }
                ],
                temperature=0.7,
                max_tokens=2500
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"生成最终提示词时出错: {e}")
            return self.accumulated_patterns

def main():
    """主函数"""
    # 配置参数
    data_directory = "./data"  # 存放JSON文件的目录
    batch_size = 10
    iterations = 3
    
    # 创建规律发现引擎
    engine = PatternDiscoveryEngine(data_directory, batch_size)
    
    try:
        # 执行规律发现
        final_patterns = engine.discover_patterns(iterations)
        
        # 保存结果
        result = {
            "final_patterns": final_patterns,
            "summary": "AI自主探索发现的谱面生成规律",
            "methodology": "采用渐进式迭代方法，通过多轮学习和验证发现规律"
        }
        
        with open("discovered_patterns.json", "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        
        print("\n规律发现完成！结果已保存到 discovered_patterns.json")
        print("\n最终规律提示词:")
        print(final_patterns)
        
    except Exception as e:
        print(f"执行过程中出错: {e}")

if __name__ == "__main__":
    main()